import os

import cv2
import h5py
import numpy as np
import torch
from torch.utils.data import Dataset


def get_path(splits, name):
    return os.path.join('data', splits, name)


class SVHNDataset(Dataset):
    def __init__(self, splits='train', transforms=None, max_length=6):
        self.splits = splits
        self.transforms = transforms
        self.max_length = max_length
        self.file = h5py.File(get_path(self.splits, 'digitStruct.mat'), 'r')
        self.name_refs = np.array(self.file['digitStruct/name']).squeeze(axis=1)
        self.bbox_refs = np.array(self.file['digitStruct/bbox']).squeeze(axis=1)

    def __deref_int(self, r):
        return int(np.array(self.file[r])) if isinstance(r, h5py.h5r.Reference) else int(r)

    def __get_all(self, idx):
        im_fn = "".join([chr(i) for i in np.array(self.file[self.name_refs[idx]]).squeeze(axis=1)])
        im = cv2.imread(get_path(self.splits, im_fn))
        b0 = self.file[self.bbox_refs[idx]]
        tt = [b0['left'], b0['top'], b0['width'], b0['height'], b0['label']]
        tt = [[self.__deref_int(j) for j in np.array(i).squeeze(axis=1)] for i in tt]
        boxes = []
        labels = []
        for x0, y0, w, h, label in zip(*tt):
            boxes.append((x0, y0, x0 + w, y0 + h))
            labels.append(label)
        label_len = len(labels)
        while len(labels) < self.max_length:
            labels.append(0)
        if len(labels) > self.max_length:
            raise ValueError("number too long, please increase max_length")
        return im, boxes, labels, label_len

    def __getitem__(self, idx):
        im, _, labels, label_len = self.__get_all(idx)
        if self.transforms is not None:
            im = self.transforms(im)
        return im, torch.tensor(labels, dtype=torch.int32), label_len

    def __len__(self):
        return self.name_refs.shape[0]

    def show_sample(self, idx):
        im_orig, boxes, labels, label_len = self.__get_all(idx)
        im = im_orig.copy()
        print(im.shape)
        label_str = "GT: " + "".join((str(i % 10) for i in labels if i != 0))
        for bbox in boxes:
            cv2.rectangle(im, bbox[0:2], bbox[2:4], (0, 255, 0), 2)
        cv2.imshow(label_str, im)
        cv2.waitKey()


def show_ds_info():
    d = SVHNDataset()
    print("{} training samples".format(len(d)))
    d.show_sample(33271)


if __name__ == '__main__':
    show_ds_info()
